import pandas
import sys
import os,re,shutil
# 获取上级目录并添加到sys.path
from pathlib import Path
parent_dir = str(Path(__file__).resolve().parent.parent)
sys.path.append(parent_dir)
sys.path.append(r"D:\ProgramData\anaconda3\envs\index-tts\index-tts")
import torch
import soundfile as sf
import logging
from datetime import datetime
import platform
from indextts.infer import IndexTTS
from pydub import AudioSegment
from pydub.silence import detect_silence
import librosa
import numpy as np
from Utils.voice_handle import adjust_speed_and_pitch,add_silence_with_pydub


class indextts_batch_generate():

    _temp= r"e:\python-knowledge\temp"

    vioce_param_map = {
        1.0: "very_low",
        2.0: "low",
        3.0: "moderate",
        4.0: "high",
        5.0: "very_high"
    }
    def __init__(self):
        self._init_tts_model()

    @classmethod
    def shorten_silence(self, input_file):
        """
        去除音频文件中的长静音片段。
        """
        # 读取音频文件
        audio = AudioSegment.from_file(input_file)
        # 检测静音片段，阈值为 -40 dBFS，最小静音时长为 100 毫秒
        silence_ranges = detect_silence(audio, min_silence_len=100, silence_thresh=-40)
        new_audio = AudioSegment.empty()
        start_point = 0
        for start, end in silence_ranges:
            # 如果静音时长超过 3 秒
            if (end - start) > 3000:
                new_audio += audio[start_point:start]
                # 只保留 1 秒的静音
                new_audio += AudioSegment.silent(duration=1000)
                start_point = end
                print(">>>>>>Silence detected ： %s"%input_file)
            else:
                new_audio += audio[start_point:end]
                start_point = end
        new_audio += audio[start_point:]
        # 导出处理后的音频文件
        return new_audio

    @classmethod
    def read_voice_data(self,seed_path,sheet_name):
        #读取音色数据excel文件,生成音色数据字典
        voice_data = pandas.read_excel(seed_path, sheet_name=sheet_name, header=0)
        keys=voice_data.columns[0] #获取第一列
        voice_data = voice_data.set_index(keys)
        voice_dict = voice_data.to_dict('index') #将数据转换为字典
        print('音色数据读取完成。')
        self.voice_dict = voice_dict
        return voice_dict

    @classmethod
    def read_text_data(self,xls_path,sheet_name):
        #读取文本数据excel文件,生成文本数据字典
        text_data = pandas.read_excel(xls_path, sheet_name=sheet_name, header=0)
        text_dict = text_data.to_dict(orient='records') #将数据转换为字典
        print('文本数据读取完成。')
        return text_dict

    def _init_tts_model(self):
        model_dir = r"D:\ProgramData\anaconda3\envs\index-tts\index-tts\checkpoints"
        self.model = IndexTTS(cfg_path=model_dir+"/config.yaml", model_dir=model_dir, device="cuda:0",is_fp16=True, use_cuda_kernel=False)


    def run_tts(self,vocice_info,text_info,save_dir,index):
        
        # Generate unique filename using timestamp
        # timestamp = datetime.now().strftime("%Y%m%d%H%M%S")
        save_path = os.path.join(save_dir, f"{index+2}.wav")
        temp_path =  os.path.join(self._temp, f"{index+2}.wav")

        logging.info("Starting inference...")

        pitch = text_info["pitch"]
        speed = text_info["speed"]
        
        # 重新获取prompt_speech_path和prompt_text
        volume = text_info["volume"]
        actor_wavs = vocice_info["prompt_speech_path"]
        wavs_list = os.listdir(actor_wavs)
        wavs_list = [os.path.join(actor_wavs, wav) for wav in wavs_list if wav.endswith(".wav") or wav.endswith(".mp3")]
        if isinstance(volume,str):
            for wav in wavs_list:
                if re.search(volume,wav):
                    prompt_speech_path = wav
                    # index_seed = re.search(r"seed_(\d+)",wav).group(1)
                    # prompt_text = self.voice_dict[int(index_seed)]["prompt_text"]
                    break
                elif re.search("default",wav):
                    default_wav = wav
            else:
                print("没有找到对应的音频文件: %s"%volume)
                prompt_speech_path = default_wav
                # prompt_text = vocice_info["prompt_text"]
        else:
            print("here")
            print(wavs_list)
            for wav in wavs_list:
                if re.search("default",wav):
                    prompt_speech_path = wav
                    break
            # prompt_text = vocice_info["prompt_text"]


        # print(text_info["text"]) 
        # print(prompt_speech_path)
        # print(prompt_text)
        # print(pitch)
        # print(speed)

        #推理    
        self.model.infer_fast(audio_prompt=prompt_speech_path, text=text_info["text"], output_path=temp_path, verbose=False)

        #处理
        audio = AudioSegment.from_wav(temp_path)
        current_dbfs = audio.dBFS # 计算当前音频的 dBFS 
        target_dbfs = -20 # 目标 dBFS 为 -20

        # 计算音量增益
        gain = target_dbfs - current_dbfs   #dbfs一般为负数
        if vocice_info["prompt_speech_path"] == "1":
            add_vol = (int(pitch)- 3)/2
        else:
            add_vol = (int(pitch)- 4)/2

        if gain > 0: # 如果当前 dBFS 小于目标 dBFS，增加音量
            audio = audio - gain*(gain/target_dbfs)  + add_vol
        elif gain < 0: # 如果当前 dBFS 大于目标 dBFS，降低音量
            if current_dbfs > -12: # 如果当前 dBFS 小于目标 dBFS，增加音量
                audio = audio + gain*(gain/-15)  + add_vol
            else:
                audio = audio + gain*((gain-2)/target_dbfs)  + add_vol
        # print(current_dbfs,audio.dBFS)
        audio.export(temp_path, format="wav")

        # 调整音调、调整语速
        speed_factor = 1+(speed-3)*0.05 # 语速因子，1 表示不调整，大于 1 表示加快，小于 1 表示减慢
        pitch_factor = (pitch-3)*0.5 # 音调因子，0表示原调 (暂时没用)
        if speed_factor != 1 and pitch_factor != 0:
            result = adjust_speed_and_pitch(temp_path,save_path,speed_factor=speed_factor,pitch_factor=pitch_factor)
            if not result:
                shutil.copy(temp_path, save_path)
        else:
            shutil.copy(temp_path, save_path)
        return save_path

    def batch_run(self,voice_dict,text_dict,save_dir,index=0):
        #需要从失败处重新开始的话，可以设置index，index=excel行号-2
        # Ensure the save directory exists
        shutil.rmtree(save_dir,ignore_errors=True)
        os.makedirs(save_dir, exist_ok=True)

        shutil.rmtree(self._temp,ignore_errors=True)   #清理临时目录
        os.makedirs(self._temp, exist_ok=True)

        # 检查角色语音是否OK，如果有错误，抛错不执行推理
        error_flag= 0
        with open("./log.log","w+",encoding="utf-8") as flog:
            for i,text_info in enumerate(text_dict):
                if i== 0:
                    flog.write(text_info["text"]+"\n")
                if text_info["seed"] not in voice_dict.keys():
                    flog.write("%d 行 ：%s 不存在 \n"%(i,text_info["seed"]))
                    error_flag = 1
        if error_flag == 1:
            raise ValueError("角色语音不存在，请检查log.log文件")

        #批量生成音频
        for i,text_info in enumerate(text_dict):
            if i >= index:    #从指定行开始
                # 执行转换               
                voice_info  = voice_dict[text_info["seed"]]
                save_path = self.run_tts(voice_info,text_info,save_dir,i)
                if i == 0 :
                    add_silence_with_pydub(save_path,save_path,silence_duration=2)


    @classmethod
    def splice_wav(cls,directory, output_file):
        # 初始化一个空的 AudioSegment 对象
        combined = AudioSegment.empty()
        # 遍历目录中的所有文件

        for filename in sorted(os.listdir(directory),key=cls.natural_keys):
            if filename.endswith('.wav'):
                file_path = os.path.join(directory, filename)
                try:
                    # 加载音频文件
                    # audio = AudioSegment.from_wav(file_path)
                    audio = cls.shorten_silence(file_path) # 去除长静音
                    

                    # 拼接音频
                    combined += audio
                except Exception as e:
                    print(f"Error processing {filename}: {e}")
        # 保存拼接后的音频文件
        if len(combined) > 0:
            combined.export(output_file, format='wav')
            print(f"Audio files concatenated and saved to {output_file}")

    @classmethod
    def atoi(cls,text):
        return int(text) if text.isdigit() else text

    @classmethod
    def natural_keys(cls,text):
        return [cls.atoi(c) for c in re.split(r'(\d+)', text)]

if __name__ == "__main__":
    convert_dir = r"E:\txt\网游之无限\网游之无限1"
    convert_list = os.listdir(convert_dir)

    indextts = indextts_batch_generate()
    seed_path = r'e:\txt\seed2.xlsx'
    sheet_name = 'sparktts'
    voice_dict = indextts.read_voice_data(seed_path,sheet_name)  # 读取音色库

    for file_name in convert_list:
        save_dir = r"e:\txt\wav"

        sheet_name = 'Sheet'
        xls_path = os.path.join(convert_dir,file_name)
        text_dict = indextts.read_text_data(xls_path,sheet_name)   # 读取文本excel

        # #存储路径
        indextts.batch_run(voice_dict,text_dict,save_dir,index=0)  #批量推理
        wavname = file_name.replace(".xlsx",".wav")
        indextts.splice_wav(save_dir, r"e:\txt\%s"%wavname)   #合成保存


    # text_info = {"text":"我认为呢","pitch":2,"speed":2,"volume":3}
    # indextts.run_tts(vocice_info=voice_dict["周星驰"],text_info=text_info,save_dir=r"E:\txt\temp",index=10)
